from maix import camera, display, nn, image, app
import time

# 初始化YOLOv8检测器
detector = nn.YOLOv8(model="/root/models/yolov8n.mud", dual_buff=True)

# 初始化摄像头和显示屏
# 使用1920*1080高分辨率拍摄以提高识别准确度
cam = camera.Camera(detector.input_width(), detector.input_height(), detector.input_format())
disp = display.Display()

# 英文标签到中文标签的映射
label_map = {
    'person': '人',
    'bicycle': '自行车',
    'car': '汽车',
    'motorcycle': '摩托车',
    'airplane': '飞机',
    'bus': '公交车',
    'train': '火车',
    'truck': '卡车',
    'boat': '船',
    'traffic light': '交通灯',
    'fire hydrant': '消防栓',
    'stop sign': '停车标志',
    'parking meter': '停车计时器',
    'bench': '长椅',
    'bird': '鸟',
    'cat': '猫',
    'dog': '狗',
    'horse': '马',
    'sheep': '羊',
    'cow': '牛',
    'elephant': '大象',
    'bear': '熊',
    'zebra': '斑马',
    'giraffe': '长颈鹿',
    'backpack': '背包',
    'umbrella': '雨伞',
    'handbag': '手提包',
    'tie': '领带',
    'suitcase': '行李箱',
    'frisbee': '飞盘',
    'skis': '滑雪板',
    'snowboard': '滑雪板',
    'sports ball': '运动球',
    'kite': '风筝',
    'baseball bat': '棒球棒',
    'baseball glove': '棒球手套',
    'skateboard': '滑板',
    'surfboard': '冲浪板',
    'tennis racket': '网球拍',
    'bottle': '瓶子',
    'wine glass': '酒杯',
    'cup': '杯子',
    'fork': '叉子',
    'knife': '刀',
    'spoon': '勺子',
    'bowl': '碗',
    'banana': '香蕉',
    'apple': '苹果',
    'sandwich': '三明治',
    'orange': '橙子',
    'broccoli': '西兰花',
    'carrot': '胡萝卜',
    'hot dog': '热狗',
    'pizza': '披萨',
    'donut': '甜甜圈',
    'cake': '蛋糕',
    'chair': '椅子',
    'couch': '沙发',
    'potted plant': '盆栽植物',
    'bed': '床',
    'dining table': '餐桌',
    'toilet': '马桶',
    'tv': '电视',
    'laptop': '笔记本电脑',
    'mouse': '鼠标',
    'remote': '遥控器',
    'keyboard': '键盘',
    'cell phone': '手机',
    'microwave': '微波炉',
    'oven': '烤箱',
    'toaster': '烤面包机',
    'sink': '水槽',
    'refrigerator': '冰箱',
    'book': '书',
    'clock': '时钟',
    'vase': '花瓶',
    'scissors': '剪刀',
    'teddy bear': '泰迪熊',
    'hair drier': '吹风机',
    'toothbrush': '牙刷'
}

while not app.need_exit():
    # 读取摄像头图像
    img = cam.read()
    
    # 进行对象检测（使用原始高分辨率图像以提高识别准确度）
    objs = detector.detect(img, conf_th=0.5, iou_th=0.45)
    
    # 绘制检测结果
    for obj in objs:
        img.draw_rect(obj.x, obj.y, obj.w, obj.h, color=image.COLOR_RED, thickness=2)
        # 使用中文标签（去除英文标签两边的空格后再查找映射）
        english_label = detector.labels[obj.class_id].strip()
        chinese_label = label_map.get(english_label, english_label)  # 如果找不到中文翻译，则使用英文标签
        label = f'{chinese_label}: {obj.score:.2f}'
        img.draw_string(obj.x, obj.y - 20, label, color=image.COLOR_RED, scale=0.5)
        print(f"检测到: {chinese_label}, 置信度: {obj.score:.2f}, 位置: ({obj.x}, {obj.y}, {obj.w}, {obj.h})")
    
    # 在显示前将图像缩放到640*480以节省内存和提高显示性能
    img = img.resize(640, 480)
    
    # 显示图像
    disp.show(img)